#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2019/2/13 9:12 PM
# @Author  : Hardy
# @File    : trie_test.py


# 字典树节点
class Node:
    def __init__(self):
        """
        字典树节点， children 用于保存字典
        """
        self.children = None


class TrieTree:
    def __init__(self, words):
        # 创建根结点
        self.root = Node()
        # 构建词典树
        self.add_keywords(words)

    def add_keywords(self, words):
        """
        添加多个词
        :param words: 词的列表
        """
        for w in words:
            self.add_word(w)

    def add_word(self, word):
        """
        添加单个词，构建词典树
        :param word: 词
        """
        node = self.root
        end = len(word) - 1
        for i in range(len(word)):
            if node.children is None:
                node.children = dict()
                # 每个children的key 是一个字，value 是一个元组
                # 这个元组第一个元素是trie树节点，第二个元素是这个字是否是一个词的最后一个字
                if i != end:
                    node.children[word[i]] = (Node(), False)
                else:
                    node.children[word[i]] = (Node(), True)

            elif word[i] not in node.children:
                if i != end:
                    node.children[word[i]] = (Node(), False)
                else:
                    node.children[word[i]] = (Node(), True)
            else:
                if i == end:
                    next_node, is_word_end = node.children[word[i]]
                    node.children[word[i]] = (next_node, True)

            node = node.children[word[i]][0]

    def get_words_by_prefix(self, prefix):
        node = self.root
        res = []
        for p in prefix:
            # 空树
            if node.children is None:
                break
            # 前缀不存在
            if p not in node.children:
                return []
            else:
                node, is_word_end = node.children[p]

        # 余下部分为空树
        if node.children is None:
            return []
        # 开始遍历剩下部分
        self.bfs(node, res, prefix)
        return res

    def bfs(self, node, out: list, k=''):
        """
        广度优先搜索, 使用尾递归

        :param node:
        :param out:
        :param k:
        """
        if node.children is None:
            return out

        for w, (next_node, is_word_end) in node.children.items():
            # print(w, is_word_end)
            if is_word_end:
                out.append(k+w)
            self.bfs(next_node, out, k + w)


if __name__ == '__main__':
    my_words = ['吃饭', '吃烤鸡', '吃烤鸭', '吃烤鸡翅', '吃']
    trie = TrieTree(my_words)

    print(trie.get_words_by_prefix('哈'))
    print(trie.get_words_by_prefix('吃'))
    print(trie.get_words_by_prefix('吃烤'))
    print(trie.get_words_by_prefix('吃烤鸡'))


